# encoding=utf-8
import argparse
import os
import time
import requests
from PIL import Image
import cv2
import numpy as np


parser = argparse.ArgumentParser()
parser.add_argument("--remote", type=str, default="http://127.0.0.1:30012/run/deeplabv3/five")
parser.add_argument("--workdir", type=str, default="workdir")
args = parser.parse_args()
if args.workdir == "" or args.workdir is None:
    print("work dir is empty")
    exit(-1)

remote_api = args.remote
work_dir = args.workdir

from_path = os.path.join(work_dir, "user_input")
to_path = os.path.join(work_dir, "user_output")
os.makedirs(to_path, exist_ok=True)


def run_remote(abs_file_path: str):
    t1 = time.perf_counter()
    f = {"file": open(abs_file_path, 'rb')}

    out = requests.post(remote_api, files=f)

    t2 = time.perf_counter()
    print("{} cost time, {}s".format(os.path.basename(abs_file_path), t2 - t1))
    return out


def process_img(image):

    colored_mask = np.zeros((image.shape[0], image.shape[1], 3), dtype=np.uint8)

    for class_id, color in class_colors.items():
        colored_mask[image == class_id] = color

    return Image.fromarray(colored_mask)


def main():
    file_names = os.listdir(from_path)
    for each_name in file_names:
        text, extension = os.path.splitext(each_name)
        abs_path = os.path.join(from_path, each_name)
        out = run_remote(abs_path)
        if out is None:
            print("处理异常或图片上传失败")
        elif out.status_code != 200:
            print("远端执行异常")
        else:
            image = cv2.imdecode(np.frombuffer(out.content, np.uint8), 0)
        result = process_img(image)
        result_path = os.path.join(to_path, f'{text}.png')
        result.save(result_path, format='PNG')
        print(f'{each_name}的结果图保存在{result_path}')


if __name__ == '__main__':
    class_colors = {
        1: (255, 0, 0),  # Red
        2: (0, 255, 0),  # Green
        3: (0, 0, 255),  # Blue
        4: (255, 255, 0),  # Yellow
        5: (255, 0, 255)  # Magenta
    }
    main()


